#ifndef JOINT_BOOSTING_H_
#define JOINT_BOOSTING_H_

#include "ZMLAlgorithm.h"
#include "DecisionStump.h"
#include <vector>
#include <set>
#include <map>

GLOBAL_NAMESPACE_BEGIN

NAMESPACE_MACHINE_LEARNING_BEGIN

class JointBoosting : public MLAlgorithm
{
private:
    struct Params
    {
        // input params
        int stepNum;
        int roundM;

        // computed params
        int dimNum;
        int classNum;
        std::vector<double> trainDataDimMaxs;
        std::vector<double> trainDataDimMins;
    };

public:
    JointBoosting();
    ~JointBoosting();

    void init(const Eigen::MatrixXd& data);
    bool commit();

    void setParams(int roundM, int stepNum=10);
    bool train(const Eigen::MatrixXd& trainData, const Eigen::VectorXi& labels);
    void evaluate(const Eigen::MatrixXd& queryData, Eigen::MatrixXd& evalMat);
    int  predict(const Eigen::RowVectorXd& oneData);

    int getLearnerSize();
    int getLearnerSize() const;
    WeakLearner* getWeakLearnerAt(int idx);
    WeakLearner* getWeakLearnerAt(int idx) const;

private:
    void destroy();
    void releaseWeakLearners(std::vector<WeakLearner*>& learners);

private:
    void updateWeights(const Eigen::MatrixXd& trainData, const WeakLearner& weakLearner, const Eigen::VectorXi& labels, Eigen::MatrixXd& weights);
    void prepareParamForWeakLearner();

private:
    Eigen::MatrixXd data_;
    std::vector<WeakLearner*> weakLearners_;
    Eigen::MatrixXd         weights_;    // N*C elements
    Eigen::MatrixXd         Hs_;         // N*C elements
    Params                  params_;
    std::vector<double>     t0c_;       // for training weak learners, C elements
    double                  totalWeights_;
};

NAMESPACE_MACHINE_LEARNING_END

GLOBAL_NAMESPACE_END

#endif//JOINT_BOOSTING_H_